from xinference.client.restful.restful_client import (
  Client,
  RESTfulChatModelHandle,
)
from vanna.base.base import VannaBase
from typing import List
import pandas as pd
from chromadb import Documents, EmbeddingFunction, Embeddings
from langchain_community.embeddings import XinferenceEmbeddings

class XinferenceEmbeddings(VannaBase):
    def __init__(self, config=None):
        VannaBase.__init__(self, config=config)
        if not config or "base_url" not in config:
            raise ValueError("config must contain at least Xinference base_url")
        self.base_url = config["base_url"]
        self.model_uid = config.get("model_uid", None)
        self.xinference_client = Client(base_url=self.base_url)
    
    def generate_embedding(self, data: str, **kwargs) -> List[float]:
        self.embeddings = self.xinference_client.get_model(self.model_uid)
        embedding = self.embeddings.create_embedding(data)
        return embedding["data"][0]["embedding"]

class Xinference_Embeddings_Function(EmbeddingFunction[Documents]):
    def __init__(self, config=None):
        VannaBase.__init__(self, config=config)

        if not config or "base_url" not in config:
            raise ValueError("config must contain at least Xinference base_url")

        self.base_url = config["base_url"]
        self.model_uid = config.get("model_uid", None)
        api_key = config.get("api_key", "not empty")
        self.xinference_client = Client(base_url=self.base_url, api_key=api_key)
        # self.embeddings = XinferenceEmbeddings(server_url=self.base_url,model_uid=self.model_uid)
        self.embeddings = self.xinference_client.get_model(self.model_uid)

    def __call__(self,input: Documents) -> Embeddings:
        # 替换输入文档中的换行符，因为换行符可能会影响嵌入性能
        input = [t.replace("\n", " ") for t in input]
        # 初始化一个空列表，用于存储所有文档的嵌入向量
        all_embeddings = []
        # 打印当前处理的文档数量
        print(f"Generating embeddings for {len(input)} documents")

        # 遍历每个文档，分别调用API生成嵌入向量
        for document in input:
            try:
                # embedding = self.embeddings.embed_query(document)
                embedding = self.embeddings.create_embedding(document)
                # 将当前文档的嵌入向量添加到列表中
                all_embeddings.append(embedding["data"][0]["embedding"])
                # all_embeddings.append(embedding)
            except Exception as e:
                raise ValueError(f"Error generating embedding for document: {e}")


        return all_embeddings

    